import os
import sys
import time

import fire

script_dir_path = os.path.dirname(__file__)
project_dir_path = os.path.join(script_dir_path, "../..")
sys.path.append(project_dir_path)

from scripts.config import preset_text_datasets_loader
from scripts.datasets.wrapping import wrap_text_dataset_for_training


def main(dataset_name: str):
    dataset = preset_text_datasets_loader[dataset_name]()
    dataset = wrap_text_dataset_for_training(dataset,
        chunk_tokens_n=2,
        batch_samples_n=5)
    for sample in dataset:
        x, y, m = sample
        print(f"{x=!r}")
        print(f"{y=!r}")
        print(f"{m=!r}")
        time.sleep(1)


if __name__ == '__main__':
    fire.Fire(main)
